package com.dyface;

import org.datavec.image.loader.NativeImageLoader;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.api.preprocessor.DataNormalization;
import org.nd4j.linalg.dataset.api.preprocessor.ImagePreProcessingScaler;

import java.io.File;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

/**
 * 人脸识别服务类
 */
public class FaceRecognitionService {
    private final FaceRecognitionModel 模型;
    private final NativeImageLoader 图片加载器;
    private final DataNormalization 数据标准化器;
    private String 当前模型路径;

    public FaceRecognitionService() {
        this.模型 = new FaceRecognitionModel();
        this.图片加载器 = new NativeImageLoader(224, 224, 3);
        this.数据标准化器 = new ImagePreProcessingScaler(0, 1);
    }

    /**
     * 初始化模型
     * @param 模型名称 模型保存名称
     * @throws IOException 如果保存失败
     */
    public void 初始化模型(String 模型名称) throws IOException {
        模型.保存模型(模型名称);
        this.当前模型路径 = 模型名称;
    }

    /**
     * 注册人脸
     * @param 图片路径 人脸图片路径
     * @param 人员ID 人员ID
     * @param 姓名 人员姓名
     * @param 模型路径 使用的模型路径
     * @throws IOException 如果图片加载失败或特征保存失败
     */
    public void 注册人脸(String 图片路径, int 人员ID, String 姓名, String 模型路径) throws IOException {
        // 如果提供了新的模型路径，则加载该模型
        if (模型路径 != null && !模型路径.equals(当前模型路径)) {
            模型.加载模型(模型路径);
            当前模型路径 = 模型路径;
        }
        
        INDArray 图片数据 = 图片加载器.asMatrix(new File(图片路径));
        数据标准化器.transform(图片数据);
        模型.注册人脸(图片数据, 人员ID, 姓名);
    }

    /**
     * 识别人脸
     * @param 图片路径 待识别的图片路径
     * @param 模型路径 使用的模型路径
     * @param 相似度阈值 识别相似度阈值（0-1之间）
     * @return 识别结果，包含ID和姓名
     * @throws IOException 如果图片加载失败
     */
    public Map<String, Object> 识别人脸(String 图片路径, String 模型路径, double 相似度阈值) throws IOException {
        // 如果提供了新的模型路径，则加载该模型
        if (模型路径 != null && !模型路径.equals(当前模型路径)) {
            模型.加载模型(模型路径);
            当前模型路径 = 模型路径;
        }

        INDArray 图片数据 = 图片加载器.asMatrix(new File(图片路径));
        数据标准化器.transform(图片数据);
        Map<String, Object> 结果 = 模型.识别人脸(图片数据, 相似度阈值);
        
        if ((double)结果.get("相似度") < 相似度阈值) {
            结果.put("ID", -1);
            结果.put("姓名", "查无此人");
        }
        
        return 结果;
    }

    /**
     * 关闭服务，释放资源
     */
    public void 关闭() {
        模型.关闭();
    }
} 